import torch
import collections


def load_model_compile(model, model_pth, device, strict=False, backend="inductor"):
    # 兼容torch1/2大版本之间的模型加载
    origin_dict = torch.load(model_pth, map_location=device)
    state_dict = collections.OrderedDict()
    # torch1_model_prefix = 'module.'
    # offset1 = len(torch1_model_prefix)
    torch2_model_prefix = '_orig_mod.'
    offset2 = len(torch2_model_prefix)
    for key, value in origin_dict.items():
        if key.startswith(torch2_model_prefix):
            if (torch.__version__).startswith('2.0'):
                model = torch.compile(model, backend=backend)
                model.load_state_dict(origin_dict, strict=strict)
            else:
                for key, value in origin_dict.items():
                    state_dict[key[offset2: len(key)]] = value
                model.load_state_dict(state_dict, strict=strict)
        else:
            if (torch.__version__).startswith('2.0'):
                model.load_state_dict(origin_dict, strict=strict)
                model = torch.compile(model, backend=backend)
            else:
                model.load_state_dict(origin_dict, strict=strict)
        break
    return model
